import network
import time
import _thread
import ulogger
from machine import Pin
from thread_shared_data import ThreadSharedData_StopOnly
from led_control import LedControl, LedStyle
import ntptime

class WifiState:
    IDLE = 0
    CONNECTING = 1
    CONNECTED = 2
    FAILURE_WAITING_RETRY = 3
    FAILURE_NEVER_RETRY = 4

class WiFiManager:
    def __init__(self, ssid: str, password: str, as_hostname: str, state_led_pin_no: int, logger: ulogger.Logger=None):
        """
        最简单的WiFi管理器
        :param ssid: WiFi名称(不可更改)
        :param password: WiFi密码(不可更改)
        :param led_pin: LED指示灯GPIO引脚
        """
        self.ssid = ssid
        self.password = password
        self.as_hostname = as_hostname
        self.failure_retry_interval = 10

        self.ntp_interval = 86400
        self.ntp_fail_intervals = (10, 60, 600, 3600)

        if logger is not None:
            self.logger = logger
        else:
            self.logger = ulogger.Logger('wifi')

        self.state_led_pin_no = state_led_pin_no

        class SharedData(ThreadSharedData_StopOnly):
            def __init__(self):
                super().__init__()
                self.wifi_state = WifiState.IDLE

        self.shared_data = SharedData()

    def enable(self):
        """启动WiFi连接线程"""
        _thread.start_new_thread(self._wifi_thread, (self.ssid, self.password))
        self.shared_data.on_thread_start()

    def disable(self):
        """停止WiFi线程并断开连接"""
        self.shared_data.request_and_wait_thread_exit()

    def get_wifi_state(self):
        with self.shared_data as shared_data:
            return shared_data.wifi_state
        
    def _set_wifi_state(self, state):
        with self.shared_data as shared_data:
            shared_data.wifi_state = state
        
    def _on_wifi_thread_exit(self, wlan: network.WLAN, led_control: LedControl):
        self.logger.info('Received Exit Request, exiting WiFi Thread...')

        try:
            wlan.disconnect()
        except:
            pass

        try:
            wlan.active(False)
        except:
            pass

        led_control.stop_led_control()

    def _wifi_thread(self, ssid, password):
        wlan = network.WLAN(network.STA_IF)

        led_control = LedControl(self.state_led_pin_no)
        led_control.start_led_control()
        
        last_ntp_success_time = None
        ntp_fail_counter = 0

        while True:
            self.logger.info('Connecting to WiFi...')

            wlan.active(True)
            wlan.config(hostname=self.as_hostname)
            wlan.connect(ssid, password)

            self._set_wifi_state(WifiState.CONNECTING)
            led_control.set_led_sytle(LedStyle.FAST_BLINK)

            # connect wait loop
            while True:
                status = wlan.status()

                if status == network.STAT_IDLE:
                    ...
                elif status == network.STAT_CONNECTING:
                    ...
                elif status == network.STAT_WRONG_PASSWORD:
                    self.logger.info('Wrong Password, exiting thread...')
                    self._set_wifi_state(WifiState.FAILURE_NEVER_RETRY)
                    self._on_wifi_thread_exit(wlan, led_control)
                    self.shared_data.quit_thread()
                    self.logger.info('Wifi thread exited.')
                elif status == network.STAT_NO_AP_FOUND:
                    self.logger.info('No AP Found')
                    wlan.disconnect()
                    self._set_wifi_state(WifiState.FAILURE_WAITING_RETRY)
                    break
                elif status == network.STAT_CONNECT_FAIL:
                    self.logger.info('Connection Failed')
                    wlan.disconnect()
                    self._set_wifi_state(WifiState.FAILURE_WAITING_RETRY)
                    break
                elif status == network.STAT_GOT_IP:
                    self.logger.info('Connected. ', wlan.ifconfig())
                    self._set_wifi_state(WifiState.CONNECTED)
                    break
                else:
                    self.logger.info(f'WTF? {status}')

                self.shared_data.check_thread_exit_request(lambda: self._on_wifi_thread_exit(wlan, led_control))
                time.sleep(0.1)

            if self.get_wifi_state() == WifiState.CONNECTED:
                # connected loop
                self.logger.info('Wifi is now in CONNECTED loop...')
                led_control.set_led_sytle(LedStyle.ON)

                while True:
                    # NTP
                    if last_ntp_success_time is None or (time.time() - last_ntp_success_time > self.ntp_interval):
                        try:
                            self.logger.info('Calling NTP...')
                            ntptime.settime()
                            self.logger.info('NTP Successfully called.')
                            last_ntp_success_time = time.time()
                            ntp_fail_counter = 0
                        except Exception as e:
                            ntp_fail_interval = self.ntp_fail_intervals[ntp_fail_counter] if ntp_fail_counter < len(self.ntp_fail_intervals) else self.ntp_fail_intervals[-1]
                            self.logger.info(f'NTP failed for the {ntp_fail_counter + 1}th time. {e}. Will retry in {ntp_fail_interval} seconds...')
                            for i in range(ntp_fail_interval):
                                self.shared_data.check_thread_exit_request(lambda: self._on_wifi_thread_exit(wlan, led_control))
                                time.sleep(1)

                            ntp_fail_counter += 1

                    # check wifi suddenly disconnected
                    if not wlan.isconnected():
                        self.logger.info('WiFi Disconnected.')
                        break

                    self.shared_data.check_thread_exit_request(lambda: self._on_wifi_thread_exit(wlan, led_control))
                    time.sleep(1)
            else:
                self.logger.info('Failure waiting...')
                led_control.set_led_sytle(LedStyle.SLOW_BLINK)

                start = time.time()
                while time.time() - start < self.failure_retry_interval:
                    self.shared_data.check_thread_exit_request(lambda: self._on_wifi_thread_exit(wlan, led_control))
                    time.sleep(0.1)



if __name__ == '__main__':
    try:
        wifi_manager = WiFiManager('XiaoYuanLian', '15818757855', 13)
        wifi_manager.enable()
    except KeyboardInterrupt:
        print('KeyboardInterrupt event!!!')
        wifi_manager.disable()